/**
 * @a https://leetcode.cn/problems/kth-largest-element-in-an-array/description/
 */

#include "common.h"

class Solution
{
public:
    int findKthLargest(vector<int> &nums, int k)
    {
        int ret = _findKthLargest(nums.begin(),nums.end(), k);
        return ret;
    }
private:
    int _findKthLargest(vector<int>::iterator begin, vector<int>::iterator end, int k)
    {
        if(begin == end - 1) return *begin;

        int x1 = *begin,x2 = *(end - 1), x3 = *(begin + rand() % (end - begin));
        // int pivot = x1 + x2 + x3 - max({x1,x2,x3}) - min({x1,x2,x3});
        int pivot = max(min(x1, x2), min(max(x1, x2), x3));
        auto left = begin, right = end - 1, curr = begin;
        while(curr <= right)
        {
            if(*curr < pivot){
                iter_swap(left++, curr++);
            }
            else if(*curr > pivot){
                iter_swap(right--,curr);
            }
            else{
                ++curr;
            }
        }
        int len_right2end = end - right - 1;
        int len_pivot = right - left + 1;
        int len_begin2left = left - begin;
        if(k <= len_right2end) return _findKthLargest(right + 1, end, k);
        else if(k <= len_right2end + len_pivot) return pivot;
        else return _findKthLargest(begin, left, k - len_pivot - len_right2end);
    }
};
class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        return findKthLargest_aux(nums.begin(), nums.end(), k);
    }
    int findKthLargest_aux(
        const vector<int>::iterator& begin,
        const vector<int>::iterator& end,
        int k)
    {
        if(begin >= end) return 0;
        auto mid = begin + ((end - begin) >> 1);
        auto left = begin, right = end - 1, cur = begin;
        int pivot = *mid;
        while(cur <= right){
            if(*cur < pivot){
                iter_swap(cur++, left++);
            }
            else if(*cur > pivot){
                iter_swap(cur, right--);
            }
            else cur++;
        }
        int len_right2end = end - right - 1;
        int len_pivot =right - left + 1;
        int len_begin2left= left - begin;
        if(k <= len_right2end) return findKthLargest_aux(right + 1, end, k);
        else if(k <= len_right2end + len_pivot) return pivot;
        else return findKthLargest_aux(begin, left, k - len_pivot - len_right2end);
    }
};